package com.allm.ai.common.strategy.impl;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.stereotype.Component;

import com.allm.ai.common.dto.ApiRequest;
import com.allm.ai.common.dto.StreamingOutput;
import com.allm.ai.common.entity.AiModel;
import com.allm.ai.common.strategy.BaseAiModelStrategy;
import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.model.chat.ChatResponse;
import com.baidubce.qianfan.model.chat.Message;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

/**
 * 百度千帆大模型策略实现
 * 支持文心大模型系列
 * 基于官方SDK: https://github.com/baidubce/bce-qianfan-sdk/tree/main/java
 */
@Slf4j
@Component
public class BaiduQianfanStrategy extends BaseAiModelStrategy {
    
    // 使用 ConcurrentHashMap 缓存不同配置的客户端
    private static final Map<String, Qianfan> CLIENT_CACHE = new ConcurrentHashMap<>();
    
    
    @Override
    public boolean supports(String serviceProvider) {
        return "Baidu".equalsIgnoreCase(serviceProvider) || 
               "Qianfan".equalsIgnoreCase(serviceProvider);
    }
    
    @Override
    public boolean supportsStreaming() {
        return true;
    }
    
    @Override
    public String callApi(AiModel aiModel, ApiRequest request) throws Exception {
        try {
            Qianfan client = getOrCreateClient(aiModel);
            
            // 构建消息列表
            List<Message> messages = buildMessages(request);
            
            // 使用官方SDK的链式调用方式
            ChatResponse response = client.chatCompletion()
                .model(aiModel.getRequestModelName()) // 使用模型名称
                .messages(messages) // 添加消息列表
                .temperature(aiModel.getTemperature() != null ? aiModel.getTemperature().doubleValue() : 0.7)
                .maxOutputTokens(aiModel.getMaxTokens() != null ? aiModel.getMaxTokens() : 2000)
                .execute(); // 发起请求
            
            // 解析响应
            return parseResponse(response);
            
        } catch (Exception e) {
            log.error("百度千帆API调用失败: {}", e.getMessage(), e);
            throw new RuntimeException("API调用失败: " + e.getMessage(), e);
        }
    }
    
    @Override
    public Flux<StreamingOutput> callApiStream(AiModel aiModel, ApiRequest request) {
        return Flux.create(sink -> {
            try {
                Qianfan client = getOrCreateClient(aiModel);
                
                // 构建消息列表
                List<Message> messages = buildMessages(request);
                
                // 异步调用流式API
                CompletableFuture.runAsync(() -> {
                    try {
                        client.chatCompletion()
                            .model(aiModel.getRequestModelName())
                            .messages(messages)
                            .temperature(aiModel.getTemperature() != null ? aiModel.getTemperature().doubleValue() : 0.7)
                            .topP(0.9d)
                            .maxOutputTokens(aiModel.getMaxTokens() != null ? aiModel.getMaxTokens() : 2000)
                            .executeStream().forEachRemaining(response -> {
                                if (response != null && response.getResult() != null && !response.getResult().trim().isEmpty()) {
                                    String content = response.getResult();
                                    if (content != null && !content.trim().isEmpty()) {
                                        // 使用基类的智能分块方法
                                        streamContentInSmartChunks(content, sink, StreamingOutput.CONTENT);
                                    }
                                }
                            });
                        
                        sink.complete();
                    } catch (Exception e) {
                        log.error("百度千帆流式API调用失败: {}", e.getMessage(), e);
                        sink.error(new RuntimeException("流式调用失败: " + e.getMessage(), e));
                    }
                });
                
            } catch (Exception e) {
                log.error("百度千帆流式调用初始化失败: {}", e.getMessage(), e);
                sink.error(new RuntimeException("流式调用初始化失败: " + e.getMessage(), e));
            }
        });
    }
    
    /**
     * 获取或创建千帆客户端（线程安全）
     * 使用双重检查锁定模式确保线程安全
     */
    private Qianfan getOrCreateClient(AiModel aiModel) {
        // 生成缓存键，基于API密钥
        String cacheKey = generateCacheKey(aiModel);
        
        // 先从缓存中获取
        Qianfan cachedClient = CLIENT_CACHE.get(cacheKey);
        if (cachedClient != null) {
            return cachedClient;
        }
        
        // 双重检查锁定
        synchronized (BaiduQianfanStrategy.class) {
            // 再次检查缓存
            cachedClient = CLIENT_CACHE.get(cacheKey);
            if (cachedClient != null) {
                return cachedClient;
            }
            
            // 创建新客户端
            Qianfan newClient = createClient(aiModel);
            CLIENT_CACHE.put(cacheKey, newClient);
            
            log.info("创建新的百度千帆客户端，缓存键: {}", cacheKey);
            return newClient;
        }
    }
    
    /**
     * 创建千帆客户端
     */
    private Qianfan createClient(AiModel aiModel) {
        // 创建千帆客户端
        Qianfan client = new Qianfan(aiModel.getAk(), aiModel.getSk());
        log.info("百度千帆客户端创建成功");
        
        return client;
    }
    
    /**
     * 生成缓存键
     * 基于API密钥生成唯一标识
     */
    private String generateCacheKey(AiModel aiModel) {
        String apiKey = aiModel.getApiKey();
        if (apiKey != null && apiKey.contains(":")) {
            // 使用API Key的前8位作为缓存键
            return "qianfan_" + apiKey.substring(0, Math.min(8, apiKey.length()));
        }
        
        // 使用环境变量的标识
        String envKey = System.getenv("QIANFAN_ACCESS_KEY");
        if (envKey != null) {
            return "qianfan_env_" + envKey.substring(0, Math.min(8, envKey.length()));
        }
        
        return "qianfan_default";
    }
    
    /**
     * 构建消息列表
     * 将ApiRequest转换为百度千帆SDK需要的格式
     */
    private List<Message> buildMessages(ApiRequest request) {
        List<Message> messages = new ArrayList<>();
        
        // 添加历史消息
        if (request.getMessages() != null) {
            for (Map<String, Object> msg : request.getMessages()) {
                Message message = new Message();
                message.setRole(convertRole((String) msg.get("role")));
                message.setContent((String) msg.get("content"));
                messages.add(message);
            }
        }
        
        // 添加当前用户消息（如果有）
        if (request.getModel() != null && !request.getModel().trim().isEmpty()) {
            // 如果ApiRequest中有用户消息，添加进去
            // 这里需要根据实际的ApiRequest结构来调整
        }
        
        return messages;
    }
    
    /**
     * 转换角色名称
     */
    private String convertRole(String role) {
        if (role == null) return "user";
        
        switch (role.toLowerCase()) {
            case "assistant":
            case "ai":
                return "assistant";
            case "system":
                return "system";
            case "user":
            default:
                return "user";
        }
    }
    
    /**
     * 解析响应
     */
    private String parseResponse(ChatResponse response) {
        if (response == null || response.getResult() == null || response.getResult().isEmpty()) {
            throw new RuntimeException("API响应为空");
        }
        
        String content = response.getResult();
        if (content == null || content.trim().isEmpty()) {
            throw new RuntimeException("API响应内容为空");
        }
        
        return content;
    }
    
    /**
     * 清理缓存（可选，用于测试或配置更新时）
     */
    public static void clearCache() {
        synchronized (BaiduQianfanStrategy.class) {
            CLIENT_CACHE.clear();
            log.info("百度千帆客户端缓存已清理");
        }
    }
    
    /**
     * 获取缓存统计信息（用于监控）
     */
    public static Map<String, Object> getCacheStats() {
        return Map.of(
            "cacheSize", CLIENT_CACHE.size(),
            "cacheKeys", CLIENT_CACHE.keySet()
        );
    }
}
